{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.plotly.v1+json": {
       "config": {
        "plotlyServerURL": "https://plot.ly"
       },
       "data": [
        {
         "hovertemplate": "step=%{x}<br>loss=%{y}<extra></extra>",
         "legendgroup": "",
         "line": {
          "color": "#636efa",
          "dash": "solid"
         },
         "marker": {
          "symbol": "circle"
         },
         "mode": "lines",
         "name": "",
         "orientation": "v",
         "showlegend": false,
         "type": "scatter",
         "x": [
          100,
          200,
          300,
          400,
          500,
          600,
          700,
          800,
          900,
          1000,
          1100,
          1200,
          1300,
          1400,
          1500,
          1600,
          1700,
          1800,
          1900,
          2000,
          2100,
          2200,
          2300,
          2400,
          2500,
          2600,
          2700,
          2800,
          2900,
          3000,
          3100,
          3200,
          3300,
          3400,
          3500,
          3600,
          3700,
          3800,
          3900,
          4000,
          4100,
          4200,
          4300,
          4400,
          4500,
          4600,
          4700,
          4800,
          4900,
          5000,
          5100,
          5200,
          5300,
          5400,
          5500,
          5600,
          5700,
          5800,
          5900,
          6000,
          6100,
          6200,
          6300,
          6400,
          6500,
          6600,
          6700,
          6800,
          6900,
          7000,
          7100,
          7200,
          7300,
          7400,
          7500,
          7600,
          7700,
          7800,
          7900,
          8000,
          8100,
          8200,
          8300,
          8400,
          8500,
          8600,
          8700,
          8800,
          8900,
          9000,
          9100,
          9200,
          9300,
          9400,
          9500,
          9600,
          9700,
          9800,
          9900,
          10000,
          10100,
          10200,
          10300,
          10400,
          10500,
          10600,
          10700,
          10800,
          10900,
          11000,
          11100,
          11200,
          11300,
          11400,
          11500,
          11600,
          11700,
          11800,
          11900,
          12000,
          12100,
          12200,
          12300,
          12400,
          12500,
          12600,
          12700,
          12800,
          12900,
          13000,
          13100,
          13200,
          13300,
          13400,
          13500,
          13600,
          13700,
          13800,
          13900,
          14000,
          14100,
          14200,
          14300,
          14400,
          14500,
          14600,
          14700,
          14800,
          14900,
          15000,
          15100,
          15200,
          15300,
          15400,
          15500,
          15600,
          15700,
          15800,
          15900,
          16000,
          16100,
          16200,
          16300,
          16400,
          16500,
          16600,
          16700,
          16800,
          16900,
          17000,
          17100,
          17200,
          17300,
          17400,
          17500,
          17600,
          17700,
          17800,
          17900,
          18000,
          18100,
          18200,
          18300,
          18400,
          18500,
          18600,
          18700,
          18800,
          18900,
          19000,
          19100,
          19200,
          19300,
          19400,
          19500,
          19600,
          19700,
          19800,
          19900
         ],
         "xaxis": "x",
         "y": [
          6.7003,
          5.0766,
          4.4222,
          4.0741,
          3.8455,
          3.679,
          3.5477,
          3.4379,
          3.3341,
          3.242,
          3.1667,
          3.0984,
          3.0517,
          3.003,
          2.9675,
          2.9338,
          2.9079,
          2.8825,
          2.8583,
          2.8387,
          2.8201,
          2.8015,
          2.7842,
          2.7649,
          2.7528,
          2.7424,
          2.7301,
          2.7168,
          2.7031,
          2.6951,
          2.6827,
          2.6745,
          2.667,
          2.6573,
          2.65,
          2.6381,
          2.6332,
          2.6283,
          2.6211,
          2.6124,
          2.6052,
          2.601,
          2.5974,
          2.5911,
          2.5829,
          2.5788,
          2.575,
          2.5688,
          2.5645,
          2.5591,
          2.5518,
          2.5467,
          2.5433,
          2.5425,
          2.5371,
          2.5327,
          2.5317,
          2.5247,
          2.5244,
          2.5175,
          2.5161,
          2.5108,
          2.5101,
          2.503,
          2.5023,
          2.4973,
          2.4963,
          2.4911,
          2.4898,
          2.4854,
          2.4833,
          2.4792,
          2.4789,
          2.4742,
          2.4718,
          2.4726,
          2.4681,
          2.4639,
          2.4598,
          2.4587,
          2.4584,
          2.4549,
          2.4535,
          2.45,
          2.4423,
          2.4387,
          2.4364,
          2.4368,
          2.4347,
          2.4332,
          2.4301,
          2.4285,
          2.4268,
          2.4283,
          2.4233,
          2.4262,
          2.4213,
          2.4199,
          2.4157,
          2.415,
          2.415,
          2.4154,
          2.4119,
          2.4123,
          2.407,
          2.4093,
          2.4058,
          2.4053,
          2.4009,
          2.4014,
          2.3975,
          2.399,
          2.3978,
          2.3982,
          2.3978,
          2.394,
          2.3929,
          2.3929,
          2.3877,
          2.3874,
          2.388,
          2.3873,
          2.3841,
          2.382,
          2.382,
          2.3799,
          2.3812,
          2.3785,
          2.3798,
          2.3756,
          2.3771,
          2.3733,
          2.3759,
          2.3741,
          2.3672,
          2.3683,
          2.3686,
          2.3653,
          2.3674,
          2.3669,
          2.3647,
          2.3625,
          2.3615,
          2.3596,
          2.362,
          2.3584,
          2.3585,
          2.3544,
          2.3549,
          2.3543,
          2.3534,
          2.3564,
          2.355,
          2.3514,
          2.3538,
          2.3509,
          2.3505,
          2.3476,
          2.3494,
          2.3477,
          2.3467,
          2.3445,
          2.3424,
          2.3432,
          2.3434,
          2.3408,
          2.3406,
          2.3397,
          2.3397,
          2.3288,
          2.3288,
          2.3274,
          2.3279,
          2.3286,
          2.325,
          2.326,
          2.327,
          2.3266,
          2.3247,
          2.3218,
          2.3249,
          2.3241,
          2.3189,
          2.3237,
          2.3219,
          2.3196,
          2.3206,
          2.3177,
          2.3186,
          2.3159,
          2.3188,
          2.3172,
          2.3156,
          2.3152,
          2.3143,
          2.316,
          2.3152,
          2.3134,
          2.3125
         ],
         "yaxis": "y"
        }
       ],
       "layout": {
        "height": 400,
        "legend": {
         "tracegroupgap": 0
        },
        "template": {
         "data": {
          "bar": [
           {
            "error_x": {
             "color": "#2a3f5f"
            },
            "error_y": {
             "color": "#2a3f5f"
            },
            "marker": {
             "line": {
              "color": "#E5ECF6",
              "width": 0.5
             },
             "pattern": {
              "fillmode": "overlay",
              "size": 10,
              "solidity": 0.2
             }
            },
            "type": "bar"
           }
          ],
          "barpolar": [
           {
            "marker": {
             "line": {
              "color": "#E5ECF6",
              "width": 0.5
             },
             "pattern": {
              "fillmode": "overlay",
              "size": 10,
              "solidity": 0.2
             }
            },
            "type": "barpolar"
           }
          ],
          "carpet": [
           {
            "aaxis": {
             "endlinecolor": "#2a3f5f",
             "gridcolor": "white",
             "linecolor": "white",
             "minorgridcolor": "white",
             "startlinecolor": "#2a3f5f"
            },
            "baxis": {
             "endlinecolor": "#2a3f5f",
             "gridcolor": "white",
             "linecolor": "white",
             "minorgridcolor": "white",
             "startlinecolor": "#2a3f5f"
            },
            "type": "carpet"
           }
          ],
          "choropleth": [
           {
            "colorbar": {
             "outlinewidth": 0,
             "ticks": ""
            },
            "type": "choropleth"
           }
          ],
          "contour": [
           {
            "colorbar": {
             "outlinewidth": 0,
             "ticks": ""
            },
            "colorscale": [
             [
              0,
              "#0d0887"
             ],
             [
              0.1111111111111111,
              "#46039f"
             ],
             [
              0.2222222222222222,
              "#7201a8"
             ],
             [
              0.3333333333333333,
              "#9c179e"
             ],
             [
              0.4444444444444444,
              "#bd3786"
             ],
             [
              0.5555555555555556,
              "#d8576b"
             ],
             [
              0.6666666666666666,
              "#ed7953"
             ],
             [
              0.7777777777777778,
              "#fb9f3a"
             ],
             [
              0.8888888888888888,
              "#fdca26"
             ],
             [
              1,
              "#f0f921"
             ]
            ],
            "type": "contour"
           }
          ],
          "contourcarpet": [
           {
            "colorbar": {
             "outlinewidth": 0,
             "ticks": ""
            },
            "type": "contourcarpet"
           }
          ],
          "heatmap": [
           {
            "colorbar": {
             "outlinewidth": 0,
             "ticks": ""
            },
            "colorscale": [
             [
              0,
              "#0d0887"
             ],
             [
              0.1111111111111111,
              "#46039f"
             ],
             [
              0.2222222222222222,
              "#7201a8"
             ],
             [
              0.3333333333333333,
              "#9c179e"
             ],
             [
              0.4444444444444444,
              "#bd3786"
             ],
             [
              0.5555555555555556,
              "#d8576b"
             ],
             [
              0.6666666666666666,
              "#ed7953"
             ],
             [
              0.7777777777777778,
              "#fb9f3a"
             ],
             [
              0.8888888888888888,
              "#fdca26"
             ],
             [
              1,
              "#f0f921"
             ]
            ],
            "type": "heatmap"
           }
          ],
          "heatmapgl": [
           {
            "colorbar": {
             "outlinewidth": 0,
             "ticks": ""
            },
            "colorscale": [
             [
              0,
              "#0d0887"
             ],
             [
              0.1111111111111111,
              "#46039f"
             ],
             [
              0.2222222222222222,
              "#7201a8"
             ],
             [
              0.3333333333333333,
              "#9c179e"
             ],
             [
              0.4444444444444444,
              "#bd3786"
             ],
             [
              0.5555555555555556,
              "#d8576b"
             ],
             [
              0.6666666666666666,
              "#ed7953"
             ],
             [
              0.7777777777777778,
              "#fb9f3a"
             ],
             [
              0.8888888888888888,
              "#fdca26"
             ],
             [
              1,
              "#f0f921"
             ]
            ],
            "type": "heatmapgl"
           }
          ],
          "histogram": [
           {
            "marker": {
             "pattern": {
              "fillmode": "overlay",
              "size": 10,
              "solidity": 0.2
             }
            },
            "type": "histogram"
           }
          ],
          "histogram2d": [
           {
            "colorbar": {
             "outlinewidth": 0,
             "ticks": ""
            },
            "colorscale": [
             [
              0,
              "#0d0887"
             ],
             [
              0.1111111111111111,
              "#46039f"
             ],
             [
              0.2222222222222222,
              "#7201a8"
             ],
             [
              0.3333333333333333,
              "#9c179e"
             ],
             [
              0.4444444444444444,
              "#bd3786"
             ],
             [
              0.5555555555555556,
              "#d8576b"
             ],
             [
              0.6666666666666666,
              "#ed7953"
             ],
             [
              0.7777777777777778,
              "#fb9f3a"
             ],
             [
              0.8888888888888888,
              "#fdca26"
             ],
             [
              1,
              "#f0f921"
             ]
            ],
            "type": "histogram2d"
           }
          ],
          "histogram2dcontour": [
           {
            "colorbar": {
             "outlinewidth": 0,
             "ticks": ""
            },
            "colorscale": [
             [
              0,
              "#0d0887"
             ],
             [
              0.1111111111111111,
              "#46039f"
             ],
             [
              0.2222222222222222,
              "#7201a8"
             ],
             [
              0.3333333333333333,
              "#9c179e"
             ],
             [
              0.4444444444444444,
              "#bd3786"
             ],
             [
              0.5555555555555556,
              "#d8576b"
             ],
             [
              0.6666666666666666,
              "#ed7953"
             ],
             [
              0.7777777777777778,
              "#fb9f3a"
             ],
             [
              0.8888888888888888,
              "#fdca26"
             ],
             [
              1,
              "#f0f921"
             ]
            ],
            "type": "histogram2dcontour"
           }
          ],
          "mesh3d": [
           {
            "colorbar": {
             "outlinewidth": 0,
             "ticks": ""
            },
            "type": "mesh3d"
           }
          ],
          "parcoords": [
           {
            "line": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "type": "parcoords"
           }
          ],
          "pie": [
           {
            "automargin": true,
            "type": "pie"
           }
          ],
          "scatter": [
           {
            "fillpattern": {
             "fillmode": "overlay",
             "size": 10,
             "solidity": 0.2
            },
            "type": "scatter"
           }
          ],
          "scatter3d": [
           {
            "line": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "marker": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "type": "scatter3d"
           }
          ],
          "scattercarpet": [
           {
            "marker": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "type": "scattercarpet"
           }
          ],
          "scattergeo": [
           {
            "marker": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "type": "scattergeo"
           }
          ],
          "scattergl": [
           {
            "marker": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "type": "scattergl"
           }
          ],
          "scattermapbox": [
           {
            "marker": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "type": "scattermapbox"
           }
          ],
          "scatterpolar": [
           {
            "marker": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "type": "scatterpolar"
           }
          ],
          "scatterpolargl": [
           {
            "marker": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "type": "scatterpolargl"
           }
          ],
          "scatterternary": [
           {
            "marker": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "type": "scatterternary"
           }
          ],
          "surface": [
           {
            "colorbar": {
             "outlinewidth": 0,
             "ticks": ""
            },
            "colorscale": [
             [
              0,
              "#0d0887"
             ],
             [
              0.1111111111111111,
              "#46039f"
             ],
             [
              0.2222222222222222,
              "#7201a8"
             ],
             [
              0.3333333333333333,
              "#9c179e"
             ],
             [
              0.4444444444444444,
              "#bd3786"
             ],
             [
              0.5555555555555556,
              "#d8576b"
             ],
             [
              0.6666666666666666,
              "#ed7953"
             ],
             [
              0.7777777777777778,
              "#fb9f3a"
             ],
             [
              0.8888888888888888,
              "#fdca26"
             ],
             [
              1,
              "#f0f921"
             ]
            ],
            "type": "surface"
           }
          ],
          "table": [
           {
            "cells": {
             "fill": {
              "color": "#EBF0F8"
             },
             "line": {
              "color": "white"
             }
            },
            "header": {
             "fill": {
              "color": "#C8D4E3"
             },
             "line": {
              "color": "white"
             }
            },
            "type": "table"
           }
          ]
         },
         "layout": {
          "annotationdefaults": {
           "arrowcolor": "#2a3f5f",
           "arrowhead": 0,
           "arrowwidth": 1
          },
          "autotypenumbers": "strict",
          "coloraxis": {
           "colorbar": {
            "outlinewidth": 0,
            "ticks": ""
           }
          },
          "colorscale": {
           "diverging": [
            [
             0,
             "#8e0152"
            ],
            [
             0.1,
             "#c51b7d"
            ],
            [
             0.2,
             "#de77ae"
            ],
            [
             0.3,
             "#f1b6da"
            ],
            [
             0.4,
             "#fde0ef"
            ],
            [
             0.5,
             "#f7f7f7"
            ],
            [
             0.6,
             "#e6f5d0"
            ],
            [
             0.7,
             "#b8e186"
            ],
            [
             0.8,
             "#7fbc41"
            ],
            [
             0.9,
             "#4d9221"
            ],
            [
             1,
             "#276419"
            ]
           ],
           "sequential": [
            [
             0,
             "#0d0887"
            ],
            [
             0.1111111111111111,
             "#46039f"
            ],
            [
             0.2222222222222222,
             "#7201a8"
            ],
            [
             0.3333333333333333,
             "#9c179e"
            ],
            [
             0.4444444444444444,
             "#bd3786"
            ],
            [
             0.5555555555555556,
             "#d8576b"
            ],
            [
             0.6666666666666666,
             "#ed7953"
            ],
            [
             0.7777777777777778,
             "#fb9f3a"
            ],
            [
             0.8888888888888888,
             "#fdca26"
            ],
            [
             1,
             "#f0f921"
            ]
           ],
           "sequentialminus": [
            [
             0,
             "#0d0887"
            ],
            [
             0.1111111111111111,
             "#46039f"
            ],
            [
             0.2222222222222222,
             "#7201a8"
            ],
            [
             0.3333333333333333,
             "#9c179e"
            ],
            [
             0.4444444444444444,
             "#bd3786"
            ],
            [
             0.5555555555555556,
             "#d8576b"
            ],
            [
             0.6666666666666666,
             "#ed7953"
            ],
            [
             0.7777777777777778,
             "#fb9f3a"
            ],
            [
             0.8888888888888888,
             "#fdca26"
            ],
            [
             1,
             "#f0f921"
            ]
           ]
          },
          "colorway": [
           "#636efa",
           "#EF553B",
           "#00cc96",
           "#ab63fa",
           "#FFA15A",
           "#19d3f3",
           "#FF6692",
           "#B6E880",
           "#FF97FF",
           "#FECB52"
          ],
          "font": {
           "color": "#2a3f5f"
          },
          "geo": {
           "bgcolor": "white",
           "lakecolor": "white",
           "landcolor": "#E5ECF6",
           "showlakes": true,
           "showland": true,
           "subunitcolor": "white"
          },
          "hoverlabel": {
           "align": "left"
          },
          "hovermode": "closest",
          "mapbox": {
           "style": "light"
          },
          "paper_bgcolor": "white",
          "plot_bgcolor": "#E5ECF6",
          "polar": {
           "angularaxis": {
            "gridcolor": "white",
            "linecolor": "white",
            "ticks": ""
           },
           "bgcolor": "#E5ECF6",
           "radialaxis": {
            "gridcolor": "white",
            "linecolor": "white",
            "ticks": ""
           }
          },
          "scene": {
           "xaxis": {
            "backgroundcolor": "#E5ECF6",
            "gridcolor": "white",
            "gridwidth": 2,
            "linecolor": "white",
            "showbackground": true,
            "ticks": "",
            "zerolinecolor": "white"
           },
           "yaxis": {
            "backgroundcolor": "#E5ECF6",
            "gridcolor": "white",
            "gridwidth": 2,
            "linecolor": "white",
            "showbackground": true,
            "ticks": "",
            "zerolinecolor": "white"
           },
           "zaxis": {
            "backgroundcolor": "#E5ECF6",
            "gridcolor": "white",
            "gridwidth": 2,
            "linecolor": "white",
            "showbackground": true,
            "ticks": "",
            "zerolinecolor": "white"
           }
          },
          "shapedefaults": {
           "line": {
            "color": "#2a3f5f"
           }
          },
          "ternary": {
           "aaxis": {
            "gridcolor": "white",
            "linecolor": "white",
            "ticks": ""
           },
           "baxis": {
            "gridcolor": "white",
            "linecolor": "white",
            "ticks": ""
           },
           "bgcolor": "#E5ECF6",
           "caxis": {
            "gridcolor": "white",
            "linecolor": "white",
            "ticks": ""
           }
          },
          "title": {
           "x": 0.05
          },
          "xaxis": {
           "automargin": true,
           "gridcolor": "white",
           "linecolor": "white",
           "ticks": "",
           "title": {
            "standoff": 15
           },
           "zerolinecolor": "white",
           "zerolinewidth": 2
          },
          "yaxis": {
           "automargin": true,
           "gridcolor": "white",
           "linecolor": "white",
           "ticks": "",
           "title": {
            "standoff": 15
           },
           "zerolinecolor": "white",
           "zerolinewidth": 2
          }
         }
        },
        "title": {
         "text": "MoE预训练损失变化"
        },
        "width": 600,
        "xaxis": {
         "anchor": "y",
         "domain": [
          0,
          1
         ],
         "title": {
          "text": "step"
         }
        },
        "yaxis": {
         "anchor": "x",
         "domain": [
          0,
          1
         ],
         "title": {
          "text": "loss"
         }
        }
       }
      },
      "text/html": [
       "<div>                            <div id=\"cd7c1609-a018-4be2-94dd-3168136750ef\" class=\"plotly-graph-div\" style=\"height:400px; width:600px;\"></div>            <script type=\"text/javascript\">                require([\"plotly\"], function(Plotly) {                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"cd7c1609-a018-4be2-94dd-3168136750ef\")) {                    Plotly.newPlot(                        \"cd7c1609-a018-4be2-94dd-3168136750ef\",                        [{\"hovertemplate\":\"step=%{x}\\u003cbr\\u003eloss=%{y}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\",\"legendgroup\":\"\",\"line\":{\"color\":\"#636efa\",\"dash\":\"solid\"},\"marker\":{\"symbol\":\"circle\"},\"mode\":\"lines\",\"name\":\"\",\"orientation\":\"v\",\"showlegend\":false,\"x\":[100,200,300,400,500,600,700,800,900,1000,1100,1200,1300,1400,1500,1600,1700,1800,1900,2000,2100,2200,2300,2400,2500,2600,2700,2800,2900,3000,3100,3200,3300,3400,3500,3600,3700,3800,3900,4000,4100,4200,4300,4400,4500,4600,4700,4800,4900,5000,5100,5200,5300,5400,5500,5600,5700,5800,5900,6000,6100,6200,6300,6400,6500,6600,6700,6800,6900,7000,7100,7200,7300,7400,7500,7600,7700,7800,7900,8000,8100,8200,8300,8400,8500,8600,8700,8800,8900,9000,9100,9200,9300,9400,9500,9600,9700,9800,9900,10000,10100,10200,10300,10400,10500,10600,10700,10800,10900,11000,11100,11200,11300,11400,11500,11600,11700,11800,11900,12000,12100,12200,12300,12400,12500,12600,12700,12800,12900,13000,13100,13200,13300,13400,13500,13600,13700,13800,13900,14000,14100,14200,14300,14400,14500,14600,14700,14800,14900,15000,15100,15200,15300,15400,15500,15600,15700,15800,15900,16000,16100,16200,16300,16400,16500,16600,16700,16800,16900,17000,17100,17200,17300,17400,17500,17600,17700,17800,17900,18000,18100,18200,18300,18400,18500,18600,18700,18800,18900,19000,19100,19200,19300,19400,19500,19600,19700,19800,19900],\"xaxis\":\"x\",\"y\":[6.7003,5.0766,4.4222,4.0741,3.8455,3.679,3.5477,3.4379,3.3341,3.242,3.1667,3.0984,3.0517,3.003,2.9675,2.9338,2.9079,2.8825,2.8583,2.8387,2.8201,2.8015,2.7842,2.7649,2.7528,2.7424,2.7301,2.7168,2.7031,2.6951,2.6827,2.6745,2.667,2.6573,2.65,2.6381,2.6332,2.6283,2.6211,2.6124,2.6052,2.601,2.5974,2.5911,2.5829,2.5788,2.575,2.5688,2.5645,2.5591,2.5518,2.5467,2.5433,2.5425,2.5371,2.5327,2.5317,2.5247,2.5244,2.5175,2.5161,2.5108,2.5101,2.503,2.5023,2.4973,2.4963,2.4911,2.4898,2.4854,2.4833,2.4792,2.4789,2.4742,2.4718,2.4726,2.4681,2.4639,2.4598,2.4587,2.4584,2.4549,2.4535,2.45,2.4423,2.4387,2.4364,2.4368,2.4347,2.4332,2.4301,2.4285,2.4268,2.4283,2.4233,2.4262,2.4213,2.4199,2.4157,2.415,2.415,2.4154,2.4119,2.4123,2.407,2.4093,2.4058,2.4053,2.4009,2.4014,2.3975,2.399,2.3978,2.3982,2.3978,2.394,2.3929,2.3929,2.3877,2.3874,2.388,2.3873,2.3841,2.382,2.382,2.3799,2.3812,2.3785,2.3798,2.3756,2.3771,2.3733,2.3759,2.3741,2.3672,2.3683,2.3686,2.3653,2.3674,2.3669,2.3647,2.3625,2.3615,2.3596,2.362,2.3584,2.3585,2.3544,2.3549,2.3543,2.3534,2.3564,2.355,2.3514,2.3538,2.3509,2.3505,2.3476,2.3494,2.3477,2.3467,2.3445,2.3424,2.3432,2.3434,2.3408,2.3406,2.3397,2.3397,2.3288,2.3288,2.3274,2.3279,2.3286,2.325,2.326,2.327,2.3266,2.3247,2.3218,2.3249,2.3241,2.3189,2.3237,2.3219,2.3196,2.3206,2.3177,2.3186,2.3159,2.3188,2.3172,2.3156,2.3152,2.3143,2.316,2.3152,2.3134,2.3125],\"yaxis\":\"y\",\"type\":\"scatter\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"title\":{\"text\":\"step\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"title\":{\"text\":\"loss\"}},\"legend\":{\"tracegroupgap\":0},\"title\":{\"text\":\"MoE预训练损失变化\"},\"width\":600,\"height\":400},                        {\"responsive\": true}                    ).then(function(){\n",
       "                            \n",
       "var gd = document.getElementById('cd7c1609-a018-4be2-94dd-3168136750ef');\n",
       "var x = new MutationObserver(function (mutations, observer) {{\n",
       "        var display = window.getComputedStyle(gd).display;\n",
       "        if (!display || display === 'none') {{\n",
       "            console.log([gd, 'removed!']);\n",
       "            Plotly.purge(gd);\n",
       "            observer.disconnect();\n",
       "        }}\n",
       "}});\n",
       "\n",
       "// Listen for the removal of the full notebook cells\n",
       "var notebookContainer = gd.closest('#notebook-container');\n",
       "if (notebookContainer) {{\n",
       "    x.observe(notebookContainer, {childList: true});\n",
       "}}\n",
       "\n",
       "// Listen for the clearing of the current output cell\n",
       "var outputEl = gd.closest('.output');\n",
       "if (outputEl) {{\n",
       "    x.observe(outputEl, {childList: true});\n",
       "}}\n",
       "\n",
       "                        })                };                });            </script>        </div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import json\n",
    "import pandas as pd\n",
    "with open('/home/user/wyf/train_moe_from_scratch/20000step_trainer_state.json' , 'r', encoding='utf-8') as f:\n",
    "    data = json.load(f)\n",
    "    log_history = data['log_history']\n",
    "    steps = []\n",
    "    losses = []\n",
    "    for i in log_history[:-1]:\n",
    "        step = i['step']\n",
    "        loss = i['loss']\n",
    "        steps.append(step)\n",
    "        losses.append(loss)\n",
    "\n",
    "import seaborn as sns\n",
    "import plotly.express as px\n",
    "\n",
    "# 创建一个DataFrame\n",
    "df = pd.DataFrame({\n",
    "    'Step': steps,\n",
    "    'Loss': losses\n",
    "})\n",
    "\n",
    "# 使用Plotly绘制散点图\n",
    "fig = px.line(df, x='Step', y='Loss', title='MoE预训练损失变化',\n",
    "                 labels={'Step': 'step', 'Loss': 'loss'},\n",
    "                 )  # 添加趋势线可选\n",
    "\n",
    "fig.update_layout(width=600, height=400)\n",
    "# 显示图表\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.randn(2, 4, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "gate = nn.Linear(10, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 4, 4])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "logits = gate(x)\n",
    "logits.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-0.4174, -0.8181,  0.2931, -1.0983],\n",
       "         [ 0.4230,  0.5691,  0.7379, -0.3439],\n",
       "         [-1.4290, -0.6250, -0.0913,  0.8918],\n",
       "         [ 0.4221,  0.4372, -0.2936, -0.2649]],\n",
       "\n",
       "        [[ 0.4224, -0.3879, -0.1051,  0.6133],\n",
       "         [-0.6835, -0.4645,  0.2799,  0.1902],\n",
       "         [ 0.0440,  0.6685, -0.5246,  0.0843],\n",
       "         [ 0.3311, -0.2890, -1.1392,  0.5881]]], grad_fn=<ViewBackward0>)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[ 0.2931, -0.4174],\n",
       "          [ 0.7379,  0.5691],\n",
       "          [ 0.8918, -0.0913],\n",
       "          [ 0.4372,  0.4221]],\n",
       " \n",
       "         [[ 0.6133,  0.4224],\n",
       "          [ 0.2799,  0.1902],\n",
       "          [ 0.6685,  0.0843],\n",
       "          [ 0.5881,  0.3311]]], grad_fn=<TopkBackward0>),\n",
       " tensor([[[2, 0],\n",
       "          [2, 1],\n",
       "          [3, 2],\n",
       "          [1, 0]],\n",
       " \n",
       "         [[3, 0],\n",
       "          [2, 3],\n",
       "          [1, 3],\n",
       "          [3, 0]]]))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "logits_topk, indices = logits.topk(2, dim=-1)\n",
    "logits_topk, indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "zeros = torch.full_like(logits, float(\"-inf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-0.4174,    -inf,  0.2931,    -inf],\n",
       "         [   -inf,  0.5691,  0.7379,    -inf],\n",
       "         [   -inf,    -inf, -0.0913,  0.8918],\n",
       "         [ 0.4221,  0.4372,    -inf,    -inf]],\n",
       "\n",
       "        [[ 0.4224,    -inf,    -inf,  0.6133],\n",
       "         [   -inf,    -inf,  0.2799,  0.1902],\n",
       "         [   -inf,  0.6685,    -inf,  0.0843],\n",
       "         [ 0.3311,    -inf,    -inf,  0.5881]]], grad_fn=<ScatterBackward0>)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sparse_logits = zeros.scatter(dim=-1, index=indices, src=logits_topk)\n",
    "sparse_logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[0.3295, 0.0000, 0.6705, 0.0000],\n",
       "         [0.0000, 0.4579, 0.5421, 0.0000],\n",
       "         [0.0000, 0.0000, 0.2723, 0.7277],\n",
       "         [0.4962, 0.5038, 0.0000, 0.0000]],\n",
       "\n",
       "        [[0.4524, 0.0000, 0.0000, 0.5476],\n",
       "         [0.0000, 0.0000, 0.5224, 0.4776],\n",
       "         [0.0000, 0.6420, 0.0000, 0.3580],\n",
       "         [0.4361, 0.0000, 0.0000, 0.5639]]], grad_fn=<SoftmaxBackward0>)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sparse_logits = F.softmax(sparse_logits, dim=-1)\n",
    "sparse_logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]],\n",
       "\n",
       "        [[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "final_outputs = torch.zeros_like(x)\n",
    "final_outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 0.8376,  1.6742,  1.5400, -2.3451,  0.9166, -0.6159,  1.2018,  1.5354,\n",
       "          -0.1724,  0.3897],\n",
       "         [-0.4154, -0.7223, -0.9576,  0.2360,  0.5441,  0.0049, -0.1311, -0.6673,\n",
       "          -1.4495,  2.5249],\n",
       "         [ 1.5304,  0.3498,  0.5821, -0.2789, -0.2131,  0.7828, -0.7165,  0.8480,\n",
       "           3.1155,  0.9457],\n",
       "         [ 0.7512, -0.5504, -0.3900,  0.9416,  0.2048, -0.8649,  0.7229,  0.0534,\n",
       "          -0.2791, -0.3948],\n",
       "         [-0.1167,  1.1454,  0.4875,  0.6351, -1.6597, -0.2928,  0.0223, -0.8946,\n",
       "           0.8921,  1.7179],\n",
       "         [ 0.1254, -0.5159,  1.1645, -1.1828,  0.8253, -0.5996, -0.9029,  0.6295,\n",
       "           0.9861,  1.1496],\n",
       "         [-0.1167,  0.0729, -0.7747,  1.1776,  0.7829,  0.6341,  0.5290, -0.1743,\n",
       "           0.5624,  1.1735],\n",
       "         [-2.2027,  0.9008,  1.6723,  1.0737,  0.5271,  0.8475, -2.7148, -2.8659,\n",
       "           0.0842,  1.7391]]),\n",
       " torch.Size([8, 10]))"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_flat = x.view(-1, x.shape[-1])\n",
    "x_flat, x_flat.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[0.3295, 0.0000, 0.6705, 0.0000],\n",
       "         [0.0000, 0.4579, 0.5421, 0.0000],\n",
       "         [0.0000, 0.0000, 0.2723, 0.7277],\n",
       "         [0.4962, 0.5038, 0.0000, 0.0000],\n",
       "         [0.4524, 0.0000, 0.0000, 0.5476],\n",
       "         [0.0000, 0.0000, 0.5224, 0.4776],\n",
       "         [0.0000, 0.6420, 0.0000, 0.3580],\n",
       "         [0.4361, 0.0000, 0.0000, 0.5639]], grad_fn=<ViewBackward0>),\n",
       " torch.Size([8, 4]))"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sparse_logits_flat = sparse_logits.view(-1, sparse_logits.shape[-1])\n",
    "sparse_logits_flat, sparse_logits_flat.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 4])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expert_mask = (indices == 0).any(-1)\n",
    "expert_mask.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ True, False, False,  True,  True, False, False,  True])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expert_mask_flat = expert_mask.view(-1)\n",
    "expert_mask_flat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 0.8376,  1.6742,  1.5400, -2.3451,  0.9166, -0.6159,  1.2018,  1.5354,\n",
       "          -0.1724,  0.3897],\n",
       "         [ 0.7512, -0.5504, -0.3900,  0.9416,  0.2048, -0.8649,  0.7229,  0.0534,\n",
       "          -0.2791, -0.3948],\n",
       "         [-0.1167,  1.1454,  0.4875,  0.6351, -1.6597, -0.2928,  0.0223, -0.8946,\n",
       "           0.8921,  1.7179],\n",
       "         [-2.2027,  0.9008,  1.6723,  1.0737,  0.5271,  0.8475, -2.7148, -2.8659,\n",
       "           0.0842,  1.7391]]),\n",
       " torch.Size([4, 10]))"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expert_input = x_flat[expert_mask_flat]\n",
    "expert_input, expert_input.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "expert = nn.Linear(10, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 0.3335, -0.4578, -0.3970, -1.2636, -0.0809, -0.0451, -0.6618, -0.3685,\n",
       "          -0.5423, -0.0130],\n",
       "         [-0.3068,  0.1800,  0.3231,  0.2886,  0.1836, -0.3680, -0.1322, -0.3956,\n",
       "          -0.1619,  0.0895],\n",
       "         [ 0.5858, -0.9985,  0.8756, -0.6312,  1.3548, -0.8566, -0.4663, -0.0338,\n",
       "           0.7733, -0.5747],\n",
       "         [ 1.7103, -0.2132,  0.8046,  0.3739,  0.7678, -0.9101, -0.9059, -0.2524,\n",
       "           0.1489,  0.6372]], grad_fn=<AddmmBackward0>),\n",
       " torch.Size([4, 10]))"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "export_output = expert(expert_input)\n",
    "export_output, export_output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.3295],\n",
       "        [0.4962],\n",
       "        [0.4524],\n",
       "        [0.4361]], grad_fn=<UnsqueezeBackward0>)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gate_scores = sparse_logits_flat[expert_mask_flat, 0].unsqueeze(1)\n",
    "gate_scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 0.1099, -0.1508, -0.1308, -0.4163, -0.0267, -0.0149, -0.2180, -0.1214,\n",
       "          -0.1787, -0.0043],\n",
       "         [-0.1522,  0.0893,  0.1603,  0.1432,  0.0911, -0.1826, -0.0656, -0.1963,\n",
       "          -0.0803,  0.0444],\n",
       "         [ 0.2651, -0.4518,  0.3962, -0.2856,  0.6130, -0.3876, -0.2110, -0.0153,\n",
       "           0.3499, -0.2600],\n",
       "         [ 0.7458, -0.0930,  0.3509,  0.1631,  0.3348, -0.3969, -0.3950, -0.1101,\n",
       "           0.0649,  0.2779]], grad_fn=<MulBackward0>),\n",
       " torch.Size([4, 10]))"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "weighted_output = export_output * gate_scores\n",
    "weighted_output, weighted_output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]],\n",
       "\n",
       "        [[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]])"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "final_outputs = torch.zeros_like(x)\n",
    "final_outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.1099, -0.1508, -0.1308, -0.4163, -0.0267, -0.0149, -0.2180,\n",
       "          -0.1214, -0.1787, -0.0043],\n",
       "         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "           0.0000,  0.0000,  0.0000],\n",
       "         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "           0.0000,  0.0000,  0.0000],\n",
       "         [-0.1522,  0.0893,  0.1603,  0.1432,  0.0911, -0.1826, -0.0656,\n",
       "          -0.1963, -0.0803,  0.0444]],\n",
       "\n",
       "        [[ 0.2651, -0.4518,  0.3962, -0.2856,  0.6130, -0.3876, -0.2110,\n",
       "          -0.0153,  0.3499, -0.2600],\n",
       "         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "           0.0000,  0.0000,  0.0000],\n",
       "         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
       "           0.0000,  0.0000,  0.0000],\n",
       "         [ 0.7458, -0.0930,  0.3509,  0.1631,  0.3348, -0.3969, -0.3950,\n",
       "          -0.1101,  0.0649,  0.2779]]], grad_fn=<IndexPutBackward0>)"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "final_outputs[expert_mask] += weighted_output\n",
    "final_outputs"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "wyf",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
